from flask import jsonify, request
import h5py
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import numpy as np
from resnet50 import ResNet


def get_similar_images():
    # 查询图片路径
    query = request.json.get("query").strip()
    # 查询出的结果路径，如果有f.savefig 则表示结果集存放路径
    result = request.json.get("result").strip()
    # 特征提取训练结果集
    index = request.json.get("index").strip()
    print(query, result)
    # query = r'D:\pythonProject5\rose\red_rose.jpg'
    # index = 'vgg_featureCNN.h5'
    # result = r'D:\pythonProject5\flower_roses'
    # read in indexed images' feature vectors and corresponding image names
    h5f = h5py.File(index, 'r')
    # feats = h5f['dataset_1'][:]
    feats = h5f['dataset_1'][:]
    print(feats)
    imgNames = h5f['dataset_2'][:]
    print(imgNames)
    h5f.close()
    # print("searching starts")
    queryImg = mpimg.imread(query)
    # plt.title("Query Image")
    plt.imshow(queryImg)
    plt.show()

    # init VGGNet16 model
    model = ResNet()
    # extract query image's feature, compute simlarity score and sort
    queryVec = model.resnet_extract_feat(query)  # 修改此处改变提取特征的网络
    print(queryVec.shape)
    print(feats.shape)
    scores = np.dot(queryVec, feats.T)
    rank_ID = np.argsort(scores)[::-1]
    rank_score = scores[rank_ID]
    # print (rank_ID)
    print(rank_score)
    # number of top retrieved images to show
    maxres = 3  # 检索出6张相似度最高的图片
    imlist = []
    for i, index in enumerate(rank_ID[0:maxres]):
        imlist.append(imgNames[index])
        print(type(imgNames[index]))
        # print("image names: " + str(imgNames[index]) + " scores: %f" % rank_score[i])
    # print("top %d images in order are: " % maxres, imlist)
    # show top #maxres retrieved result one by one
    imlist2 = []
    for i, im in enumerate(imlist):
        imagePath = result + "/" + str(im, 'utf-8')
        imlist2.append(imagePath);
    #     image = mpimg.imread(imagePath)
    #     plt.title("search output %d" % (i + 1))
    #     plt.imshow(np.uint8(image))
    #     f = plt.gcf()  # 获取当前图像
    #     f.savefig(r'D:\pythonProject5\result\{}.jpg'.format(i), dpi=100)
    # f.clear()  # 释放内存
    # plt.show()
    return jsonify(
            {"code": 1002, "msg": "账号或密码不能为空！", "list": imlist2}
        )
